#!/usr/bin/env python3

import copy
import json
from glob import glob
import os
import sys
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../utils')
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../operators')
import wrf_operators as wrf
from utils import cli, cmd_args, parse_config, run, get_value, has_key

args = cmd_args()

config = parse_config(args.config_json)

def run_wps(config):
	wrf.config_wps(args.work_root, args.wps_root, args.geog_root, config, args)
	wrf.run_wps_geogrid(args.work_root, args.wps_root, config, args)
	wrf.run_wps_ungrib_metgrid(args.work_root, args.wps_root, args.bkg_root, config, args)

def run_spinup(config):
	wrf.config_wrf(args.work_root, args.wrf_root, config, args, tag='spinup')
	wrf.run_real(args.work_root, args.wrf_root, config, args, tag='spinup')
	wrf.run_wrf(args.work_root, args.wrf_root, config, args, tag='spinup')

def run_ose(exp, config, blacklist=False, whitelist=False):
	work_dir = f'{args.work_root}/{exp}'
	wrf.config_wrf(args.work_root, args.wrf_root, config, args, tag=exp)
	wrf.run_real(args.work_root, args.wrf_root, config, args, tag=exp)
	if os.path.isdir(f'{args.work_root}/spinup'):
		start_time_str = config['custom']['start_time'].format('YYYY-MM-DD_HH:mm:ss')
		wrfinput = f'{args.work_root}/spinup/wrf/wrfout_d01_{start_time_str}'
		if not os.path.isfile(wrfinput):
			cli.error(f'File {wrfinput} does not exist! Check spinup run!')
		run(f'cp --remove-destination {wrfinput} {work_dir}/wrf/{os.path.basename(wrfinput).replace("wrfout", "wrfinput")}')
		run(f'ln -sf {work_dir}/wrf/{os.path.basename(wrfinput).replace("wrfout", "wrfinput")} {work_dir}/wrf/wrfinput_d01')

	if has_key(config, ('custom', 'blacklist')):
		wrf.config_wrfda(args.work_root, args.wrfda_root, config, args, tag=exp)
		wrf.run_wrfda_obsproc(args.work_root, args.wrfda_root, config, args, tag=exp, littler_root=args.littler_root, blacklist=blacklist)
		wrf.run_wrfda_3dvar(args.work_root, args.wrfda_root, config, args, tag=exp)
		wrf.run_wrfda_update_bc(args.work_root, args.wrfda_root, config, args, tag=exp)
	elif has_key(config, ('custom', 'whitelist')):
		wrf.config_wrfda(args.work_root, args.wrfda_root, config, args, tag=exp)
		wrf.run_wrfda_obsproc(args.work_root, args.wrfda_root, config, args, tag=exp, littler_root=args.littler_root, whitelist=whitelist)
		wrf.run_wrfda_3dvar(args.work_root, args.wrfda_root, config, args, tag=exp)
		wrf.run_wrfda_update_bc(args.work_root, args.wrfda_root, config, args, tag=exp)
		

	wrf.run_wrf(args.work_root, args.wrf_root, config, args, tag=exp)
	wrf.run_upp(args.work_root, args.upp_root, config, args, tag=exp)
	wrf.run_met(args.work_root, args.met_root, config, args, tag=exp)

spinup_hours = get_value(config['custom'], 'spinup_hours', 0)

wps_config = copy.deepcopy(config)
if spinup_hours > 0:
	wps_config['custom']['start_time'] = config['custom']['start_time'].subtract(hours=spinup_hours)
	wps_config['custom']['forecast_hours'] += spinup_hours
run_wps(wps_config)

# Spinup.
if spinup_hours > 0:
	spinup_config = copy.deepcopy(config)
	spinup_config['custom']['start_time'] = config['custom']['start_time'].subtract(hours=spinup_hours)
	spinup_config['custom']['forecast_hours'] = spinup_hours
	spinup_config['custom']['end_time'] = spinup_config['custom']['start_time'].add(hours=spinup_hours)
	run_spinup(spinup_config)

# Run reference experiment without DA.
run_ose('ref', config, blacklist=True)

# Run sensitivity experiment with DA.
run_ose('sens', config)

# Merge two met json files.
try:
	json_file = os.path.basename(glob(f'{args.work_root}/ref/met/d01_*.json')[0])
	d1 = json.load(open(f'{args.work_root}/ref/met/{json_file}', 'r'))
	d2 = json.load(open(f'{args.work_root}/sens/met/{json_file}', 'r'))
	json.dump({'ref': d1,'sens': d2}, open(f'{args.work_root}/result.json', 'w'))
except:
	cli.warning('There is no MET result in JSON format.')
